# -*- coding: utf-8 -*-  
'''
训练tplinker网络

@author: luoyi
Created on 2021年6月28日
'''
import glob
import os
import sys
#    取项目根目录
ROOT_PATH = os.path.abspath(os.path.dirname(__file__)).split('knowledge_graph')[0]
ROOT_PATH = ROOT_PATH + "knowledge_graph"
sys.path.append(ROOT_PATH)

import data.dataset_baidu as ds_baidu
import data.dataset_tfrecord_baidu as dstf_baidu
import utils.conf as conf
import utils.relationships as rel
import utils.dictionaries as dicts

from models.tplinker.nets import TPLinkerNet


#    初始化字典，关系库
dicts.load_dict_from_pkl()
dicts.auto_append_off()
rel.load_rel_id_from_pkl()


#    百度训练集，验证集
batch_size = conf.DATASET_BAIDU.get_batch_size()
#    取多少个文件
max_file_idx = conf.DATASET_BAIDU.get_max_file_idx()
#    取文件总数
tplinker_dataset_path=conf.DATASET_BAIDU.get_train_tplinker_dataset_path()
if (max_file_idx < 0): files = glob.glob(os.path.dirname(tplinker_dataset_path) + '/*.tfrecord')
else: files = [tplinker_dataset_path.format(i) for i in range(max_file_idx + 1)]
#    每个文件记录数
count_file = conf.DATASET_BAIDU.get_record_count()
#    计算迭代步数
sample_total = 167609 if (max_file_idx < 0) else (max_file_idx + 1) * count_file
steps_per_epoch = sample_total // batch_size
#    数据解析方式
tfrecord_reader=dstf_baidu.TPLinkerTFRecordReader(max_sen_len=conf.TPLINKER.get_max_sentence_len(), 
                                                  rel_size=len(rel.id_rel), )

db_train = tfrecord_reader.tensor_db(tplinker_dataset_path=conf.DATASET_BAIDU.get_val_tplinker_dataset_path(), 
                                     max_file_idx=0, 
                                     batch_size=batch_size, 
                                     epochs=conf.DATASET_BAIDU.get_epochs(), 
                                     shuffle_buffer_rate=conf.DATASET_BAIDU.get_shuffle_buffer_rate(), 
                                     tfrecord_buffer_rate=conf.DATASET_BAIDU.get_tfrecord_buffer_rate(),
                                    )
db_val = tfrecord_reader.tensor_db(tplinker_dataset_path=conf.DATASET_BAIDU.get_val_tplinker_dataset_path(), 
                                   max_file_idx=0, 
                                   batch_size=batch_size, 
                                   epochs=conf.DATASET_BAIDU.get_epochs(), 
                                   shuffle_buffer_rate=conf.DATASET_BAIDU.get_shuffle_buffer_rate(), 
                                   tfrecord_buffer_rate=conf.DATASET_BAIDU.get_tfrecord_buffer_rate(),
                                   )

#    tplinker网络
tplinker = TPLinkerNet(name='tplinker', 
                       #    bert相关配置
                       vocab_size=dicts.dict_size(), 
                       n_block=conf.BERT.get_n_block(), 
                       n_head=conf.BERT.get_n_head_attention(), 
                       d_model=conf.BERT.get_d_model(), 
                       f_model=conf.BERT.get_f_model(), 
                       dropout_rate=conf.BERT.get_dropout_rate(), 
                       #    tplinker相关配置
                       max_sen_len=conf.TPLINKER.get_max_sentence_len(), 
                       rel_size=len(rel.id_rel), 
                       batch_size=conf.DATASET_BAIDU.get_batch_size(), 
                       #    模型训练相关配置
                       loss_lamda_ner=conf.TPLINKER.get_loss_lamda_ner(), 
                       loss_lamda_re=conf.TPLINKER.get_loss_lamda_re(), 
                       learning_rate=conf.TPLINKER.get_learning_rate(), 
                       input_shape=(None, 2, conf.TPLINKER.get_max_sentence_len()), 
                       auto_assembling=True, 
                       is_build=True)
tplinker.show_info()


#    喂数据
tplinker.train_tensor_db(db_train=db_train, 
                         db_val=db_val, 
                         steps_per_epoch=steps_per_epoch, 
                         batch_size=batch_size, 
                         epochs=conf.DATASET_BAIDU.get_epochs(), 
                         auto_save_weights_after_traind=True, auto_save_weights_dir=conf.TPLINKER.get_model_save_weights_path(), 
                         auto_learning_rate_schedule=True, 
                         auto_tensorboard=True, auto_tensorboard_dir=conf.TPLINKER.get_tensorboard_dir_path())




